import numpy as np

from implement.layers.basic.model import Model
from implement.layers.conv.conv2d import Conv2d
from implement.layers.basic.affine import Affine
from utils.file_opera import get_file
from utils.functions_collect import dropout, pooling, reshape, relu

'''
VGG16（Visual Geometry Group 16）是一个用于图像分类和对象识别任务的卷积神经网络（CNN）模型。它由牛津大学的Visual Geometry Group开发，是VGG系列中的一个重要成员。VGG16由16个卷积层和3个全连接层组成，因此得名。

VGG16的主要用途包括：
图像分类： VGG16经过在大规模图像数据集上的预训练，可以识别图像中的多个类别，例如ImageNet数据集中的1000个类别。
特征提取： 由于VGG16具有深度的卷积结构，它可以用作特征提取器。在许多计算机视觉任务中，可以使用VGG16的中间卷积层的输出作为图像特征。
迁移学习： 预训练的VGG16模型可以用于迁移学习。通过冻结模型的底层权重，可以将其应用于新任务，而不会破坏在原始任务上学到的特征。
研究和教育： VGG16是一个经典的深度学习模型，对于学术研究和教育来说具有重要意义。通过研究VGG16，人们可以了解卷积神经网络的基本结构和设计原理。
总体而言，VGG16在计算机视觉领域中广泛用于各种图像处理任务，成为深度学习模型设计的重要参考之一。
'''

class VGG16(Model):
    WEIGHTS_PATH = 'https://github.com/koki0702/dezero-models/releases/download/v0.1/vgg16.npz'

    def __init__(self, pretrained=False):
        super().__init__()
        # 定义卷积层
        self.conv1_1 = Conv2d(64, kernel_size=3, stride=1, pad=1)
        self.conv1_2 = Conv2d(64, kernel_size=3, stride=1, pad=1)
        self.conv2_1 = Conv2d(128, kernel_size=3, stride=1, pad=1)
        self.conv2_2 = Conv2d(128, kernel_size=3, stride=1, pad=1)
        self.conv3_1 = Conv2d(256, kernel_size=3, stride=1, pad=1)
        self.conv3_2 = Conv2d(256, kernel_size=3, stride=1, pad=1)
        self.conv3_3 = Conv2d(256, kernel_size=3, stride=1, pad=1)
        self.conv4_1 = Conv2d(512, kernel_size=3, stride=1, pad=1)
        self.conv4_2 = Conv2d(512, kernel_size=3, stride=1, pad=1)
        self.conv4_3 = Conv2d(512, kernel_size=3, stride=1, pad=1)
        self.conv5_1 = Conv2d(512, kernel_size=3, stride=1, pad=1)
        self.conv5_2 = Conv2d(512, kernel_size=3, stride=1, pad=1)
        self.conv5_3 = Conv2d(512, kernel_size=3, stride=1, pad=1)

        # 定义全连接层
        self.fc6 = Affine(4096)
        self.fc7 = Affine(4096)
        self.fc8 = Affine(1000)

        # 如果指定使用预训练权重，则加载
        if pretrained:
            weights_path = get_file(VGG16.WEIGHTS_PATH)
            self.load_weights(weights_path)

    def forward(self, x):
        # 卷积层、ReLU激活和池化
        x = relu(self.conv1_1(x))
        x = relu(self.conv1_2(x))
        x = pooling(x, 2, 2)
        x = relu(self.conv2_1(x))
        x = relu(self.conv2_2(x))
        x = pooling(x, 2, 2)
        x = relu(self.conv3_1(x))
        x = relu(self.conv3_2(x))
        x = relu(self.conv3_3(x))
        x = pooling(x, 2, 2)
        x = relu(self.conv4_1(x))
        x = relu(self.conv4_2(x))
        x = relu(self.conv4_3(x))
        x = pooling(x, 2, 2)
        x = relu(self.conv5_1(x))
        x = relu(self.conv5_2(x))
        x = relu(self.conv5_3(x))
        x = pooling(x, 2, 2)

        # 展平并对全连接层应用dropout
        x = reshape(x, (x.shape[0], -1))
        x = dropout(relu(self.fc6(x)))
        x = dropout(relu(self.fc7(x)))

        # 最后的全连接层
        x = self.fc8(x)
        return x

    @staticmethod
    def preprocess(image, size=(224, 224), dtype=np.float32):
        # 对输入图像进行预处理，以符合VGG16的要求
        image = image.convert('RGB')
        if size:
            image = image.resize(size)
        image = np.asarray(image, dtype=dtype)
        image = image[:, :, ::-1]
        image -= np.array([103.939, 116.779, 123.68], dtype=dtype)
        image = image.transpose((2, 0, 1))
        return image


